// Package traceid pkg/traceid/traceid.go
package traceid

import (
	"context"

	"github.com/gin-gonic/gin"
)

// TraceIDKey 用于存储TraceID的键
const TraceIDKey = "trace_id"

// GetTraceID 从上下文获取TraceID
func GetTraceID(c *gin.Context) string {
	traceID, exists := c.Get(TraceIDKey)
	if !exists {
		return ""
	}

	traceIDStr, ok := traceID.(string)
	if !ok {
		return ""
	}

	return traceIDStr
}

// ContextWithTraceID 创建带有TraceID的上下文
func ContextWithTraceID(ctx context.Context, traceID string) context.Context {
	return context.WithValue(ctx, TraceIDKey, traceID)
}

// GetTraceIDFromContext 从上下文获取TraceID
func GetTraceIDFromContext(ctx context.Context) string {
	traceID, ok := ctx.Value(TraceIDKey).(string)
	if !ok {
		return ""
	}
	return traceID
}
